from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad,unpad
import base64
from Crypto.Random import random
from Crypto.PublicKey import RSA
import Crypto.Signature.pkcs1_15 as DSA
from Crypto.Hash import SHA
from Crypto.Signature import DSS
import rsa
from Crypto.Signature import PKCS1_v1_5 as Signature_pkcs1_v1_5
from rsa.key import PrivateKey
from Crypto.Hash import MD5


# 读入 加密 和重写文件
def open_and_read(filepath):
    file = open( filepath, 'r+',encoding = 'utf-8' )
    text = file.read()
    file.close()
    return text

def encrypt(filepath):
    plaintext =open_and_read(filepath)
    key = input('input key(16 bytes): ')
    cipher = AES.new(key.encode('utf-8'),AES.MODE_ECB)
    paded = pad(plaintext.encode('utf-8'),16)
    ciphertext = cipher.encrypt(paded)
    encodestr = base64.b64encode(ciphertext)
    encodetext = encodestr.decode('utf-8')
    print(encodetext)
    file = open( filepath, 'w+',encoding = 'utf-8' )
    file.write(encodetext)
    file.close()
    return encodetext

def decrypt(filepath):
    ciphertext = open_and_read(filepath)
    key = input("key: ")
    ciphertext = base64.b64decode(ciphertext)
    cipher = AES.new(key.encode('utf-8'),AES.MODE_ECB)
    decrypted = unpad(bytes(cipher.decrypt(ciphertext)),16)
    text = decrypted.decode('utf-8')
    file = open(filepath, 'w+',encoding = 'utf-8' )
    file.write(text)
    print(text)
    file.close()
    return text
def sign(filepath,privkey_path):
    message = open_and_read(filepath)
    with open(privkey_path,"rb") as f:
        key = f.read()

    rsakey = RSA.importKey(key)
    signer = Signature_pkcs1_v1_5.new(rsakey)# 构建签名
    digest = SHA.new()
    digest.update(bytes(message,encoding='utf-8')) # 使用SHA算法获得摘要
    
    sign = signer.sign(digest) # 使用摘要签名
    signature = base64.b64encode(sign) #序列化

    # with open(filepath,'a+',encoding = 'utf-8') as f:
    #     f.write('\n'+str(signature))
    with open('signature.txt','wb') as f:
        f.write(signature)
    print('签名：' + str(signature) +'已写入病历文件')

    return signature

def verify(signature_path,text_path,publickey_path):
    message = open_and_read(text_path)
    key = open_and_read(publickey_path)

    rsakey = RSA.importKey(key)
    verifier = Signature_pkcs1_v1_5.new(rsakey)
    digest = SHA.new()
    digest.update(bytes(message,'utf-8'))

    with open(signature_path,'rb') as f:
        signature = f.read()
    is_verify = verifier.verify(digest,base64.decodebytes(signature))#对比解码后的签名和原文的摘要已确认
    print(is_verify)
    return is_verify

def md5_hash(filepath):
    text = open_and_read(filepath)

    digest = MD5.new()
    digest.update(bytes(text,encoding='utf-8'))
    hash = digest.hexdigest()
    print('生成数字摘要：'+ str(hash))

    return hash




if __name__=='__main__':
    filepath = input('文件路径:')
    print('1.加密 2.解密 3.签名 4.验证签名')
    choice = input()
    if choice == '1':
        encrypt(filepath)
    elif choice == '2':
        decrypt(filepath)
    elif choice == '3':
        privkey=input('私钥文件路径：')
        sign(filepath,privkey)
        md5_hash(filepath)
    elif choice == '4':
        signature = input('签名文件路径:')
        publickey=input('公钥文件路径：')
        verify(signature,filepath,publickey)
    else:
        print('error')
